from typing import Dict, List
from .tools import (
    LiteratureSearchTool,
    TextGenerationTool,
    GrammarCheckTool,
    FormattingTool
)
from .planner import OutlinePlanner
from .memory import ContextManager

class PaperWritingAgent:
    def __init__(self):
        self.literature_tool = LiteratureSearchTool()
        self.text_generator = TextGenerationTool()
        self.grammar_checker = GrammarCheckTool()
        self.formatter = FormattingTool()
        self.planner = OutlinePlanner()
        self.context_manager = ContextManager()
    
    def analyze_topic(self, topic: str, requirements: str) -> str:
        """分析论文主题"""
        analysis = self.planner.analyze_topic(topic, requirements)
        self.context_manager.add_context("topic_analysis", analysis)
        return str(analysis)
    
    def generate_outline(self, topic_analysis: str) -> str:
        """生成论文大纲"""
        outline = self.planner.generate_outline(eval(topic_analysis))
        self.context_manager.add_context("outline", outline)
        return str(outline)
    
    def write_section(self, section_name: str, outline: str) -> str:
        """生成特定章节的内容"""
        # 获取上下文信息
        topic_analysis = self.context_manager.get_context("topic_analysis")
        
        # 根据不同章节设置不同的提示
        section_prompts = {
            "摘要": "请根据以下大纲编写一个简洁明了的学术论文摘要，突出研究目的、方法、结果和结论：",
            "引言": "请根据以下大纲编写论文引言，需要包含研究背景、问题陈述、研究目的和意义：",
            "文献综述": "请根据以下大纲编写文献综述，需要系统地回顾和评述相关研究：",
            "方法": "请根据以下大纲详细描述研究方法，包括研究设计、数据收集和分析方法：",
            "结果": "请根据以下大纲呈现研究结果，使用清晰的语言描述发现：",
            "讨论": "请根据以下大纲展开讨论，解释结果的含义，并与现有研究进行对比：",
            "结论": "请根据以下大纲总结研究的主要发现，指出局限性和未来研究方向："
        }
        
        prompt = f"""
        {section_prompts.get(section_name, "请根据以下大纲编写相应章节：")}
        
        大纲内容：
        {outline}
        
        主题分析：
        {str(topic_analysis)}
        """
        
        # 生成内容
        content = self.text_generator.generate(prompt)
        
        # 缓存生成的内容
        self.context_manager.cache_result(f"section_{section_name}", content)
        
        return content
    
    def review_and_polish(self, content: str) -> str:
        """审查和润色内容"""
        # 首先进行语法检查
        corrected_content = self.grammar_checker.check_and_correct(content)
        
        # 使用 AI 模型进行内容优化
        polish_prompt = f"""
        请对以下学术论文内容进行润色和改进，注意：
        1. 提高语言的学术性和专业性
        2. 确保逻辑连贯性
        3. 增强表达的准确性
        4. 改善段落之间的过渡
        
        原文：
        {corrected_content}
        """
        
        polished_content = self.text_generator.generate(polish_prompt)
        return polished_content
    
    def generate_references(self, cited_papers: List[Dict]) -> str:
        """生成参考文献列表"""
        references = []
        for paper in cited_papers:
            # 按照 APA 格式生成引用
            authors = ", ".join(paper.get("authors", []))
            year = paper.get("year", "n.d.")
            title = paper.get("title", "")
            url = paper.get("url", "")
            
            reference = f"{authors} ({year}). {title}. {url}"
            references.append(reference)
        
        return "\n".join(references)
    
    def generate_full_paper(self, outline: str, max_words: int = 20000) -> str:
        """根据提纲生成完整论文"""
        try:
            # 获取上下文信息
            topic_analysis = self.context_manager.get_context("topic_analysis")
            
            # 分配每个部分的大致字数
            word_distribution = {
                "摘要": int(max_words * 0.05),  # 5%
                "引言": int(max_words * 0.15),  # 15%
                "文献综述": int(max_words * 0.25),  # 25%
                "方法": int(max_words * 0.20),  # 20%
                "结果": int(max_words * 0.15),  # 15%
                "讨论": int(max_words * 0.15),  # 15%
                "结论": int(max_words * 0.05)   # 5%
            }
            
            # 生成各个部分
            paper_sections = {}
            for section, word_limit in word_distribution.items():
                prompt = f"""
                请根据以下大纲生成论文的{section}部分，字数限制在{word_limit}字左右。

                大纲内容：
                {outline}

                主题分析：
                {str(topic_analysis)}

                要求：
                1. 保持学术写作风格
                2. 确保逻辑连贯性
                3. 适当引用相关研究
                4. 使用专业术语
                5. 遵循学术写作规范
                """
                
                content = self.text_generator.generate(prompt)
                paper_sections[section] = content
                
                # 缓存每个部分的内容
                self.context_manager.cache_result(f"section_{section}", content)
            
            # 整合所有部分
            paper_content = {
                "title": topic_analysis.get("title", "Research Paper"),
                "authors": topic_analysis.get("authors", []),
                "abstract": paper_sections["摘要"],
                "introduction": paper_sections["引言"],
                "literature_review": paper_sections["文献综述"],
                "methods": paper_sections["方法"],
                "results": paper_sections["结果"],
                "discussion": paper_sections["讨论"],
                "conclusion": paper_sections["结论"]
            }
            
            # 格式化论文
            formatted_paper = self.formatter.format_paper(paper_content)
            return formatted_paper
            
        except Exception as e:
            print(f"论文生成出错: {str(e)}")
            return ""